import ast
from openai import OpenAI
import re

# ---------- LLM 初始化 ----------
api_key = "sk-34f5c792513c423f90b404d28b070f1f"
base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

class LargeLanguageModel:
    def __init__(self, model: str, api_key: str, base_url: str):
        self.model = model
        self.client = OpenAI(api_key=api_key, base_url=base_url)

    def chat(self, query: str) -> str:
        completion = self.client.chat.completions.create(
            model=self.model,
            messages=[{'role': 'user', 'content': query}]
        )
        return completion.choices[0].message.content


llm = LargeLanguageModel(model="qwen-turbo", api_key=api_key, base_url=base_url)


# ---------- SQL 候选生成 ----------
def generate_sql_candidates_with_llm(query: str, metadata: dict, data_info: dict, num_candidates: int = 3):
    """
    使用大语言模型生成 SQL 候选，结合表结构信息（保留并标注字段中的单位）
    :param query: 自然语言查询
    :param metadata: {'tags': [...], 'rating': int}
    :param data_info: 表结构信息（格式如 {'表名': {'columns': [...]}, ...}）
    :param num_candidates: 候选 SQL 数量
    :return: SQL 字符串列表
    """

    # 示例 SQL 用于 few-shot 提示
    few_shots = [
        {
            "query": "返回所有用户的平均流量使用量",
            "metadata": {"tags": ["project"], "rating": 100},
            "sql": "SELECT AVG(月流量使用（GB）) FROM 用户信息表"
        },
        {
            "query": "返回每个地市中流量超出费用最高的前3个用户编码",
            "metadata": {"tags": ["project", "group", "order", "limit"], "rating": 300},
            "sql": "SELECT 用户编码 FROM 用户信息表 GROUP BY 地市 ORDER BY MAX(流量超出费用) DESC LIMIT 3"
        },
        {
            "query": "返回每个地市的平均上传速度",
            "metadata": {"tags": ["project", "group", "aggregation"], "rating": 200},
            "sql": "SELECT 地市, AVG(上传速度) AS 平均上传速度 FROM 用户信息表 GROUP BY 地市"
        },
    #     {
    #         "query": "返回所有用户的平均使用时间",
    #         "metadata": {"tags": ["project", "aggregation"], "rating": 150},
    #         "sql": "SELECT AVG(使用时间) AS 平均使用时间 FROM 用户信息表"
    #     },
    #     {
    #         "query": "返回所有用户的总流量使用量",
    #         "metadata": {"tags": ["project", "aggregation"], "rating": 120},
    #         "sql": "SELECT SUM(月流量使用（GB）) AS 总流量使用量 FROM 用户信息表"
    #     },
    #     {
    #         "query": "返回每个地市的总上传流量",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 180},
    #         "sql": "SELECT 地市, SUM(上传流量) AS 总上传流量 FROM 用户信息表 GROUP BY 地市"
    #     },
    #     {
    #         "query": "返回用户数量最多的地市",
    #         "metadata": {"tags": ["project", "group", "order", "limit"], "rating": 250},
    #         "sql": "SELECT 地市 FROM 用户信息表 GROUP BY 地市 ORDER BY COUNT(用户编码) DESC LIMIT 1"
    #     },
    #     {
    #         "query": "返回每个地市的最大流量使用量",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 220},
    #         "sql": "SELECT 地市, MAX(月流量使用（GB）) AS 最大流量使用量 FROM 用户信息表 GROUP BY 地市"
    #     },
    #     {
    #         "query": "返回每个地市的最小上传速度",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 200},
    #         "sql": "SELECT 地市, MIN(上传速度) AS 最小上传速度 FROM 用户信息表 GROUP BY 地市"
    #     },
    #     {
    #         "query": "返回流量使用量超过100GB的用户数量",
    #         "metadata": {"tags": ["project", "filter", "aggregation"], "rating": 170},
    #         "sql": "SELECT COUNT(用户编码) AS 用户数量 FROM 用户信息表 WHERE 月流量使用（GB） > 100"
    #     },
    #     {
    #         "query": "返回每个地市的平均流量超出费用",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 210},
    #         "sql": "SELECT 地市, AVG(流量超出费用) AS 平均流量超出费用 FROM 用户信息表 GROUP BY 地市"
    #     },
    #     {
    #         "query": "返回所有用户的总流量超出费用",
    #         "metadata": {"tags": ["project", "aggregation"], "rating": 130},
    #         "sql": "SELECT SUM(流量超出费用) AS 总流量超出费用 FROM 用户信息表"
    #     },
    #     {
    #         "query": "返回每个地市的用户数量",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 190},
    #         "sql": "SELECT 地市, COUNT(用户编码) AS 用户数量 FROM 用户信息表 GROUP BY 地市"
    #     },
    #     {
    #         "query": "返回上传速度大于50Mbps的用户数量",
    #         "metadata": {"tags": ["project", "filter", "aggregation"], "rating": 160},
    #         "sql": "SELECT COUNT(用户编码) AS 用户数量 FROM 用户信息表 WHERE 上传速度 > 50"
    #     },
    #     {
    #         "query": "返回每个地市的最大上传速度",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 230},
    #         "sql": "SELECT 地市, MAX(平均上传速度) AS 最大上传速度 FROM 用户信息表 GROUP BY 地市"
    #     },
    #     {
    #         "query": "返回上传速度大于10Mbps的用户的平均下载速度",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 230},
    #         "sql": "SELECT '符合条件的平均下载速度', AVG(平均下载速度) AS 平均下载速度 FROM 用户信息表 WHERE 平均上传速度 > 10"
    #     },
    #     {
    #         "query": "返回所有地市的总用户数",
    #         "metadata": {"tags": ["project", "group", "aggregation"], "rating": 230},
    #         "sql": "SELECT '总用户数', COUNT(*) AS 总用户数 FROM 用户信息表"
    #     }
    ]

    # Helper：提取单位
    def extract_column_unit(col_name: str) -> str:
        match = re.search(r'[（(](.*?)[）)]', col_name)
        return match.group(1) if match else None

    # 表结构信息格式化
    data_info_lines = []

    # 1. 单表结构
    if "table_name" in data_info and "columns" in data_info:
        table_name = data_info["table_name"]
        column_names = list(data_info["columns"].keys())
        processed_columns = []
        for col in column_names:
            unit = extract_column_unit(col)
            if unit:
                col_name_only = re.sub(r'[（(].*?[）)]', '', col)
                processed_columns.append(f"{col_name_only}（单位：{unit}）")
            else:
                processed_columns.append(col)
        data_info_lines.append(f"表名: {table_name}，字段: {', '.join(processed_columns)}")

    # 2. 多表结构
    else:
        for table_name, table_meta in data_info.items():
            if isinstance(table_meta, dict):
                columns = table_meta.get("columns", [])
                if isinstance(columns, list):
                    processed_columns = []
                    for col in columns:
                        unit = extract_column_unit(col)
                        if unit:
                            col_name_only = re.sub(r'[（(].*?[）)]', '', col)
                            processed_columns.append(f"{col_name_only}（单位：{unit}）")
                        else:
                            processed_columns.append(col)
                    data_info_lines.append(f"表名: {table_name}，字段: {', '.join(processed_columns)}")

    formatted_data_info = "\n".join(data_info_lines)

    # Prompt 模板
    sql_prompt = f"""
以下是数据库的表结构信息：
{formatted_data_info}

请注意：字段名中包含的单位（如 GB、元）是字段名的一部分，请不要省略或删除括号中的单位。生成 SQL 时应使用字段的完整名字。

以下是自然语言查询与 SQL 的转换示例：
"""

    for ex in few_shots:
        sql_prompt += f"""
### 示例
自然语言查询: {ex['query']}
Metadata: tags={ex['metadata']['tags']}, rating={ex['metadata']['rating']}
SQL: {ex['sql']}
"""

    sql_prompt += f"""

### 当前任务
自然语言查询: {query}
Metadata: tags={metadata.get('tags')}, rating={metadata.get('rating')}
请你结合上述表结构，生成 {num_candidates} 个等价但结构不同的 SQL 查询，返回格式如下（不要使用 markdown 代码块）：
[
  "SQL语句1",
  "SQL语句2",
  ...
]
"""

    # 调用 LLM
    response = llm.chat(sql_prompt)
    print("candidates: ", response)
    return response


if __name__ == "__main__":
    query = "查询所有套餐费用低于100的用户姓名"
    metadata = {
        "tags": ["filter"],
        "rating": 100
    }

    data_info = {
        "用户信息表": {
            "columns": ["用户编码", "姓名", "套餐费用", "月流量使用", "流量超出费用", "地市"]
        },
        "订单表": {
            "columns": ["订单ID", "用户编码", "金额", "下单时间"]
        }
    }

    sql_candidates = generate_sql_candidates_with_llm(query, metadata, data_info, num_candidates=3)
    print("\n✅ 生成的 SQL 候选：")
    for sql in sql_candidates:
        print("-", sql)
